import torch
import gradio as gr
from torchvision import transforms
from model import MLP

# 加载训练好的模型参数
model = MLP()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 定义数据转换函数
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 定义 Gradio 预测函数
def predict_digit(image):
    image = transform(image)
    image = image.unsqueeze(0)
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output.data, 1)
    return str(predicted.item())

# 创建 Gradio 接口
interface = gr.Interface(
    fn=predict_digit,
    inputs="sketchpad",
    outputs="text",
    title="手写数字识别",
    description="使用 PyTorch 的 MLP 进行手写数字识别"
)

# 启动 Gradio 接口
interface.launch()
